#include "conn_track.h"
#ifdef BRIDGE_MODE_SUPPORT
///////////////
pcap_output_t* Pcap_onu = NULL;
pcap_output_t* Pcap_all = NULL;
/////////////

#define TCP_HEADER_LEN(offset) ((offset)*4)

#define SUN_PATH "/tmp/ctk_sock"
TAILQ_HEAD(ConnTrack_head, conn_entry)
connTrackHead;

static pthread_mutex_t connLock;

struct {
    int fd;
    struct ifreq ifr;
    struct sockaddr_ll ll_dest;
} Output_if;

static conn_base_info_t ConnBase_info;

#define CONN_LOCK_INIT() pthread_mutex_init(&connLock, NULL)
#define CONN_LOCK() pthread_mutex_lock(&connLock)
#define CONN_UNLOCK() pthread_mutex_unlock(&connLock)

void conn_base_set_ifname(char* ifname)
{
    strcpy(ConnBase_info.ifname, ifname);
}
void conn_base_set_ethhdr(char* dmac, char* smac, uint16_t proto)
{
    memcpy(ConnBase_info.eth_hdr.h_dest, dmac, ETH_ALEN);
    memcpy(ConnBase_info.eth_hdr.h_source, smac, ETH_ALEN);
    ConnBase_info.eth_hdr.h_proto = htons(proto);
}
void conn_base_set_vlanhdr()
{
}
void conn_base_set_ppphdr()
{
}
void conn_base_set_myip(uint32_t saddr)
{
    ConnBase_info.my_ip = saddr;
}
void conn_base_set_nameserver(uint32_t addr)
{
    ConnBase_info.dns_server = addr;
}
uint32_t conn_dns_server()
{
    return ConnBase_info.dns_server;
}
inline static void pkt2tuple(Packet_t* ipkt, struct _tuple* tuple)
{
    bzero(tuple, sizeof(struct _tuple));
    tuple->dest_ip = ipkt->ip_hdr->daddr;
    tuple->src_ip = ipkt->ip_hdr->saddr;
    if (ipkt->udp_hdr) {
        tuple->dest_port = ipkt->udp_hdr->dest;
        tuple->src_port = ipkt->udp_hdr->source;
        tuple->proto = IPPROTO_UDP;
    } else if (ipkt->tcp_hdr) {
        tuple->dest_port = ipkt->tcp_hdr->dest;
        tuple->src_port = ipkt->tcp_hdr->source;
        tuple->proto = IPPROTO_TCP;
    } else if (ipkt->icmp_hdr) {
        tuple->icmpId = ipkt->icmp_hdr->icmp_id;
        tuple->proto = IPPROTO_ICMP;
    }
}
static uint16_t _ip_checksum(uint8_t* ptr, int size)
{
    int cksum = 0;
    int index = 0;

    *(ptr + 10) = 0;
    *(ptr + 11) = 0;

    if (size % 2 != 0)
        return 0;

    while (index < size) {
        cksum += *(ptr + index + 1);
        cksum += *(ptr + index) << 8;

        index += 2;
    }

    while (cksum > 0xffff) {
        cksum = (cksum >> 16) + (cksum & 0xffff);
    }
    return ~cksum;
}
static uint16_t _icmp_check_sum(uint16_t* addr, int nleft)
{
    unsigned sum = 0;
    while (nleft > 1) {
        sum += *addr++;
        nleft -= 2;
    }

    if (nleft == 1) {
        if (BB_LITTLE_ENDIAN) {
            sum += *(uint16_t*)addr;
        } else
            sum += *(uint16_t*)addr << 8;
    }

    sum = (sum >> 16) + (sum & 0xffff);
    sum += (sum >> 16);

    return (uint16_t)~sum;
}

static uint16_t _tcp_udp_checksum(uint32_t saddr, uint32_t daddr, uint16_t len,
    uint8_t protocol, const uint16_t* data)
{
    uint16_t proto = 0;
    uint32_t cksum = 0;
    int count = 0;

    proto |= protocol;
    cksum = htons((saddr >> 16)) + htons((saddr & 0xffff)) + htons((daddr >> 16)) + htons((daddr & 0xffff)) + htons(len) + (proto);
    count = ntohs(len);
    // log_buf(LOG_TRACE, data, count);
    while (count > 1) {
        cksum += htons(*data++);
        count -= 2;
    }
    if (count) {
        cksum += htons(*data);
    }
    while (cksum >> 16) {
        cksum = (cksum & 0xffff) + (cksum >> 16);
    }
    return htons((uint16_t)(~cksum));
}

static int _packet_checksum(Packet_t* packet)
{
    struct iphdr* ip_hdr = packet->ip_hdr;

    assert(ip_hdr != NULL);

    if (packet->ip_hdr) {
        ip_hdr->check = 0;
        ip_hdr->check = htons(_ip_checksum((uint8_t*)ip_hdr, ip_hdr->ihl * 4));
    }

    if (packet->icmp_hdr) {
        struct icmp* icmp_hdr = packet->icmp_hdr;
        icmp_hdr->icmp_cksum = 0;
        icmp_hdr->icmp_cksum = _icmp_check_sum((uint16_t*)icmp_hdr, sizeof(struct icmp));
        return 1;
    } else if (packet->udp_hdr) {
        struct udphdr* udp_hdr = packet->udp_hdr;
        udp_hdr->check = 0;
        udp_hdr->check = _tcp_udp_checksum(ip_hdr->saddr, ip_hdr->daddr, udp_hdr->len,
            ip_hdr->protocol, (uint16_t*)udp_hdr);
    }
    if (packet->tcp_hdr) {
        struct tcphdr* tcp_hdr = packet->tcp_hdr;
        uint16_t tcp_all_len = ntohs(ip_hdr->tot_len) - ip_hdr->ihl * 4;
        tcp_hdr->check = 0;
        tcp_hdr->check = _tcp_udp_checksum(ip_hdr->saddr, ip_hdr->daddr,
            htons(tcp_all_len),
            ip_hdr->protocol, (uint16_t*)tcp_hdr);
    }
    return 1;
}

static struct conn_entry* _conn_track_find_entry(struct _tuple* tuple, int direct)
{
    struct conn_entry* pdata = NULL;
    TAILQ_FOREACH(pdata, &connTrackHead, tailq)
    {
#if 0
        log_trace("+++++++++++++++++++ pkt tuple ++++++++++++++++++++++++");
        log_buf(LOG_TRACE, tuple, sizeof(struct _tuple));
        log_trace("-------------------- conn tuple -----------------------");
        log_buf(LOG_TRACE, &pdata->tuple[direct], sizeof(struct _tuple));
        log_trace("###########################################");
#endif
        if (memcmp(tuple, &pdata->tuple[direct], sizeof(struct _tuple)) == 0) {
            // log_trace("pdata:%p", pdata);
            break;
        }
    }
    return pdata;
}
static void _conn_track_forward(Packet_t* ipkt)
{
    struct _tuple tuple = { 0 };

    pkt2tuple(ipkt, &tuple);

    CONN_LOCK();
    struct conn_entry* entry = _conn_track_find_entry(&tuple, CONN_TUPLE_REPLY);
    if (!entry) {
        CONN_UNLOCK();
        return;
    }
    pcap_output(Pcap_onu, ipkt);
    if (0 > write(entry->fd, ipkt->data, ipkt->data_len)) {
        log_syserr("conn track write error.");
        CONN_UNLOCK();
        return;
    }
    CONN_UNLOCK();
}
struct conn_entry* conn_track_entry_new()
{
    struct conn_entry* entry = malloc(sizeof(struct conn_entry));
    if (!entry) {
        log_syserr("conn entry new error.");
        return NULL;
    }
    bzero(entry, sizeof(struct conn_entry));
    return entry;
}

void conn_track_free_entry(struct conn_entry* entry)
{
    CONN_LOCK();
    TAILQ_REMOVE(&connTrackHead, entry, tailq);
    free(entry);
    CONN_UNLOCK();
}
void conn_track_add_entry(struct conn_entry* entry)
{
    CONN_LOCK();
    TAILQ_INSERT_TAIL(&connTrackHead, entry, tailq);
    CONN_UNLOCK();
}
void conn_track_remove_entry(struct conn_entry* entry)
{
    CONN_LOCK();
    TAILQ_REMOVE(&connTrackHead, entry, tailq);
    conn_track_free_entry(entry);
    CONN_UNLOCK();
}
int conn_track_hook(Packet_t* ipkt)
{
    if (!ipkt->ip_hdr)
        return 0;

    pcap_output(Pcap_all, ipkt);

    _conn_track_forward(ipkt);
    return 0;
}
static void _reverse_tuple(struct conn_entry* entry)
{
    struct _tuple *direct, *reply;

    direct = &entry->tuple[CONN_TUPLE_DIRECT];
    reply = &entry->tuple[CONN_TUPLE_REPLY];

    reply->dest_ip = direct->src_ip;
    reply->src_ip = direct->dest_ip;
    reply->dest_port = direct->src_port;
    reply->src_port = direct->dest_port;
    reply->proto = direct->proto;
    reply->icmpId = direct->icmpId;
}
static void _fill_ip_header(struct iphdr* ip_hdr, struct _tuple* tuple)
{
    ip_hdr->version = 4;
    ip_hdr->ihl = 5;
    ip_hdr->tos = 0;
    ip_hdr->tot_len = sizeof(struct iphdr);
    ip_hdr->id = (uint16_t)get_random();
    ip_hdr->frag_off = 0;
    ip_hdr->ttl = 64;
    ip_hdr->protocol = tuple->proto;
    ip_hdr->check = 0;
    ip_hdr->saddr = ConnBase_info.my_ip;
    ip_hdr->daddr = tuple->dest_ip;
}
static void _fill_udp_header(struct udphdr* udp_hdr, struct _tuple* tuple)
{
    udp_hdr->dest = (tuple->dest_port);
    udp_hdr->source = (tuple->src_port);
    udp_hdr->len = 0;
    udp_hdr->check = 0;
}
static void _fill_tcp_header(struct tcphdr* tcp_hdr, struct _tuple* tuple, struct conn_entry* entry)
{
    struct conn_tcp_info* tcp_info = &entry->tcp_info;
    tcp_hdr->source = tuple->src_port;
    tcp_hdr->dest = tuple->dest_port;
    tcp_hdr->seq = 0;
    tcp_hdr->ack_seq = 0;
    tcp_hdr->doff = 5;
    tcp_hdr->res1 = 0;
    tcp_hdr->cwr = 0;
    tcp_hdr->ece = 0;
    tcp_hdr->urg = 0;
    tcp_hdr->ack = 0;
    tcp_hdr->psh = 0;
    tcp_hdr->rst = 0;
    tcp_hdr->syn = 0;
    tcp_hdr->fin = 0;
    tcp_hdr->window = 0xffff;
    tcp_hdr->check = 0;
    tcp_hdr->urg_ptr = 0;

    switch (tcp_info->status) {
    case TCP_SYN_SENT:
        tcp_hdr->syn = 1;
        tcp_info->sent_seq = get_random();
        tcp_hdr->seq = htonl(tcp_info->sent_seq);
        break;
    case TCP_SYN_RECV:
        tcp_hdr->ack = 1;
        tcp_info->sent_seq = htonl(tcp_info->sent_seq + 1);
        tcp_hdr->seq = tcp_info->sent_seq;
        tcp_info->rcv_seq = htonl(tcp_info->rcv_seq + 1);
        tcp_hdr->ack_seq = tcp_info->rcv_seq;
        break;
    }

}
static void _fill_packet_header(Packet_t* pkt, struct _tuple* tuple, struct conn_entry* entry)
{
    char* next_hdr = (char*)pkt->buf;
    pkt->eth_hdr = (struct ethhdr*)next_hdr;
    memcpy(pkt->eth_hdr, &ConnBase_info.eth_hdr, sizeof(struct ethhdr));
    next_hdr = (char*)pkt->eth_hdr + sizeof(struct ethhdr);
    pkt->len += sizeof(struct ethhdr);

    if (0) {
        pkt->ppp_hdr = (struct ppphdr*)next_hdr;
        memcpy(pkt->ppp_hdr, &ConnBase_info.ppp_hdr, sizeof(struct ppphdr));
        next_hdr = (char*)pkt->ppp_hdr + sizeof(struct ppphdr);
        pkt->len += sizeof(struct ppphdr);
    }

    pkt->ip_hdr = (struct iphdr*)next_hdr;
    _fill_ip_header(pkt->ip_hdr, tuple);
    next_hdr = (char*)pkt->ip_hdr + sizeof(struct iphdr);
    pkt->len += sizeof(struct iphdr);

    switch (tuple->proto) {
    case IPPROTO_ICMP:
        pkt->data = next_hdr;
        pkt->icmp_hdr = (struct icmp*)next_hdr;
        break;
    case IPPROTO_UDP:
        pkt->udp_hdr = (struct udphdr*)next_hdr;
        _fill_udp_header(pkt->udp_hdr, tuple);
        pkt->data = (char*)pkt->udp_hdr + sizeof(struct udphdr);
        pkt->len += sizeof(struct udphdr);
        break;
    case IPPROTO_TCP:
        pkt->tcp_hdr = (struct tcphdr*)next_hdr;
        _fill_tcp_header(pkt->tcp_hdr, tuple, entry);
        pkt->data = (char*)pkt->tcp_hdr + TCP_HEADER_LEN(pkt->tcp_hdr->doff);
        pkt->len += TCP_HEADER_LEN(pkt->tcp_hdr->doff);
        break;
    }
}
static void _fill_packet_data(Packet_t* pkt, char* buf, int len)
{
    if (buf)
        memcpy(pkt->data, buf, len);

    pkt->len += len;
    if (pkt->icmp_hdr) {
        pkt->ip_hdr->tot_len = htons(pkt->ip_hdr->tot_len + len);
    } else if (pkt->udp_hdr) {
        pkt->ip_hdr->tot_len = htons(pkt->ip_hdr->tot_len + len + sizeof(struct udphdr));
        pkt->udp_hdr->len = htons(len + sizeof(struct udphdr));
    } else if (pkt->tcp_hdr) {
        struct tcphdr* tcp_hdr = pkt->tcp_hdr;
        pkt->ip_hdr->tot_len = htons(pkt->ip_hdr->tot_len + len + TCP_HEADER_LEN(tcp_hdr->doff));
    }
}
static int _send_packet(Packet_t* packet)
{
    int len;

    PRINT_PACKET(packet);
    pcap_output(Pcap_onu, packet);
    if (Output_if.fd == 0) {
        Output_if.fd = socket_packet_new();
        strcpy(Output_if.ifr.ifr_name, ConnBase_info.ifname);
        ioctl(Output_if.fd, SIOCGIFINDEX, &Output_if.ifr);
        Output_if.ll_dest.sll_family = AF_PACKET;
        Output_if.ll_dest.sll_ifindex = Output_if.ifr.ifr_ifindex;
        Output_if.ll_dest.sll_halen = htons(ETH_HLEN);
    }

    len = sendto(Output_if.fd, packet->buf, packet->len, 0, (struct sockaddr*)&Output_if.ll_dest, sizeof(Output_if.ll_dest));
    if (len < 0) {
        log_syserr("net stack send error.");
    }
    // close(Output_fd);
    return len;
}
void conn_track_output(struct conn_entry* entry, char* buf, int len)
{
    Packet_t packet;
    bzero(&packet, sizeof(packet));
    _fill_packet_header(&packet, &entry->tuple[CONN_TUPLE_DIRECT], entry);
    _fill_packet_data(&packet, buf, len);
    _packet_checksum(&packet);
    _send_packet(&packet);
}
void* conn_track_task(void* arg)
{
    int fd = 0, cfd = 0, max_fd = 0;
    int ret = 0;
    char buf[1500];
    fd_set rdfds, Rdfds;
    struct timeval tmout;
    struct sockaddr_un un;
    struct conn_entry *entry, *pdata;

    fd = socket(PF_UNIX, SOCK_STREAM, 0);
    if (fd < 0) {
        log_syserr("socket create error.");
        return NULL;
    }
    un.sun_family = AF_UNIX;
    strcpy(un.sun_path, SUN_PATH);
    unlink(un.sun_path);
    if (bind(fd, (struct sockaddr*)&un, sizeof(un)) < 0) {
        log_syserr("conn track bind error.");
        close(fd);
        return NULL;
    }
    if (listen(fd, 100) < 0) {
        log_syserr("conn track listen error.");
        close(fd);
        return NULL;
    }
    max_fd = fd;
    FD_ZERO(&rdfds);
    FD_ZERO(&Rdfds);
    FD_SET(fd, &Rdfds);
    while (1) {
        tmout.tv_sec = 1;
        tmout.tv_usec = 0;
        memcpy(&rdfds, &Rdfds, sizeof(fd_set));
        ret = select(max_fd + 1, &rdfds, NULL, NULL, &tmout);
        if (ret <= 0) {
            continue;
        }
        if (FD_ISSET(fd, &rdfds)) {
            cfd = accept(fd, NULL, NULL);
            if (cfd < 0) {
                log_syserr("conn track accept error.");
                return NULL;
            }
            entry = conn_track_entry_new();
            entry->fd = cfd;
            if (read(cfd, &entry->tuple[CONN_TUPLE_DIRECT], sizeof(struct _tuple)) < 0) {
                log_syserr("conn read error.");
                conn_track_free_entry(entry);
                continue;
            }
            // log_buf(LOG_DEBUG, &entry->tuple[CONN_TUPLE_DIRECT], sizeof(struct _tuple));
            _reverse_tuple(entry);
            conn_track_add_entry(entry);
            FD_SET(cfd, &Rdfds);
            max_fd = MAX_VAL(max_fd, cfd);
            log_trace("max_fd:%d, cfd:%d", max_fd, cfd);
            continue;
        }
        TAILQ_FOREACH(pdata, &connTrackHead, tailq)
        {
            if (FD_ISSET(pdata->fd, &rdfds)) {
                ret = read(pdata->fd, buf, sizeof(buf));
                log_trace("read fd(%d) ret:%d", fd, ret);
                if (ret < 0) {
                    log_syserr("read error.");
                    continue;
                } else if (ret > 0) {
                    conn_track_output(pdata, buf, ret);
                } else if (ret == 0) {
                    FD_CLR(pdata->fd, &Rdfds);
                    close(pdata->fd);
                    conn_track_free_entry(pdata);
                }
            }
        }
    }
}

int conn_tcp_connect(conn_track* ct, int fd)
{
    int ret = 0;
    char buf[1500] = { 0 };
    struct tcphdr* tcp_hdr;

    bzero(buf, sizeof(buf));
    CONN_LOCK();
    struct conn_entry* entry = _conn_track_find_entry(ct, CONN_TUPLE_DIRECT);
    entry->refcnt++;
    CONN_UNLOCK();
    if (!entry) {
        log_trace("entry is NULL");
        return 0;
    }
    conn_entry_set_status(entry, TCP_SYN_SENT);
    conn_track_output(entry, NULL, 0);
    ret = read(fd, buf, sizeof(buf));
    log_trace("ret:%d", ret);
    if (ret < 0) {
        entry->refcnt--;
        log_syserr("tcp connect error.");
        return 0;
    }
    tcp_hdr = (struct tcphdr*)buf;
    if (tcp_hdr->ack && tcp_hdr->syn) {
        entry->tcp_info.rcv_seq = ntohl(tcp_hdr->seq);
        conn_entry_set_status(entry, TCP_SYN_RECV);
        conn_track_output(entry, NULL, 0);
    }
    entry->refcnt--;
    return 1;
}

int conn_track_regist(conn_track* ct)
{
    int fd;
    struct _tuple tuple;
    struct sockaddr_un un;

    fd = socket(PF_UNIX, SOCK_STREAM, 0);
    un.sun_family = AF_UNIX;
    strcpy(un.sun_path, SUN_PATH);
    if (connect(fd, (struct sockaddr*)&un, sizeof(un)) < 0) {
        log_syserr("conn track conn error.");
        close(fd);
        return -1;
    }
    memset(&tuple, 0, sizeof(tuple));
    tuple.dest_ip = ct->dest_ip;
    tuple.dest_port = ct->dest_port;
    tuple.src_ip = ConnBase_info.my_ip;
    if (ct->proto != IPPROTO_ICMP) {
        tuple.src_port = (ct->src_port) ? ct->src_port : port_random(0);
    }
    tuple.proto = ct->proto;
    tuple.icmpId = ct->icmpId;
    // log_buf(LOG_DEBUG, &tuple, sizeof(struct _tuple));
    if (write(fd, &tuple, sizeof(tuple)) < 0) {
        log_syserr("conn track write error.");
        close(fd);
        return -1;
    }
    memcpy(ct, &tuple, sizeof(tuple));
    return fd;
}

int conn_track_init(int argc, char** argv)
{
    Pcap_onu = pcap_output_new("onu.pcap", 0, 0);
    pcap_set_sync(Pcap_onu);
    // pcap_set_disable(Pcap_onu);

    Pcap_all = pcap_output_new("onu_all.pcap", 0, 0);
    pcap_set_sync(Pcap_all);
    // pcap_set_disable(Pcap_all);

    memset(&Output_if, 0, sizeof(Output_if));
    TAILQ_INIT(&connTrackHead);
    CONN_LOCK_INIT();

    pthread_t tid;
    pthread_create(&tid, NULL, conn_track_task, NULL);
    pthread_detach(tid);
    return 0;
}
#endif // BRIDGE_MODE_SUPPORT